/*
 * @Author: _LJP 2634716776@qq.com
 * @Date: 2023-12-29 14:34:21
 * @LastEditors: _LJP 2634716776@qq.com
 * @LastEditTime: 2024-01-06 16:00:35
 * @FilePath: /my_-linux/Pro24/AVL/NewAVL/AVLTree.hpp
 * @Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
 */
#pragma once

#include<iostream>

#include<assert.h>

using namespace std;

template<class K,class V>
struct AVLTreeNode{
    pair<K, V> _kv;
    AVLTreeNode<K, V> *_left;
    AVLTreeNode<K, V> *_right;
    AVLTreeNode<K, V> *_parent;
    int _bf ;
    
    AVLTreeNode(const pair<K,V> &kv)
    :_kv(kv)
    ,_left(nullptr)
    ,_right(nullptr)
    ,_parent(nullptr)
    ,_bf(0)
   
    {}

};

template<class K,class V>
class AVLTree{
    typedef AVLTreeNode<K,V> Node;
    public:
        int getHeight()
        {
            return getHeight(_root);
        }

    bool Insert(const std::pair<K,V>& kv){
        if (_root == nullptr)
        {
            _root = new Node(kv);
            return true;
        }

        Node *parent = nullptr;
        Node *cur = _root;
        while (cur)
        {
            if (cur->_kv.first < kv.first)
            {
                parent = cur;
                cur = cur->_right;
            }
            else if (cur->_kv.first > kv.first)
            {
                parent = cur;
                cur = cur->_left;
            }
            else
            {
                return false;
            }
        }

        cur = new Node(kv);
        if (parent->_kv.first < kv.first)
        {
            parent->_right = cur;
        }
        else
        {
            parent->_left = cur;
        }

        cur->_parent = parent;

        // ... 控制平衡
        // 更新平衡因子
        while (parent)
        {
            if (cur == parent->_left)
            {
                parent->_bf--;
            }
            else // if (cur == parent->_right)
            {
                parent->_bf++;
            }

            if (parent->_bf == 0)
            {
                // 更新结束
                break;
            }
            else if (parent->_bf == 1 || parent->_bf == -1)
            {
                // 继续往上更新
                cur = parent;
                parent = parent->_parent;
            }
            else if (parent->_bf == 2 || parent->_bf == -2)
            {
                // 子树不平衡了，需要旋转
                if (parent->_bf == 2 && cur->_bf == 1)
                {
                    RotateL(parent);
                }
                else if (parent->_bf == -2 && cur->_bf == -1)
                {
                    RotateR(parent);
                }
                else if (parent->_bf == 2 && cur->_bf == -1)
                {
                    RotateRL(parent);
                }
                else if (parent->_bf == -2 && cur->_bf == 1)
                {
                    RotateLR(parent);
                }

                break;
            }
            else
            {
                assert(false);
            }
        }

        return true;

    }//Insert函数反括号

    bool IsBalance(){
        return IsBalance(_root);
    }
    bool IsBalance(Node *root){
        if(root == nullptr){
            return true;
        }

        int leftH = getHeight(root->_left); // the leftH is subtree height for left;
        int rightH = getHeight(root->_right);

        if(rightH - leftH != root->_bf){
            cout << "平衡因子异常:" << root->_kv.first << "->" << root->_bf << endl;
            return false;
        }//

        return abs(rightH - leftH) < 2 
        && IsBalance(root->_left) 
        && IsBalance(root->_right);
    }

    void InOrder()
    {
        _InOrder(_root);
    }

    protected:

    
    int getHeight(Node *root){
        if(root == nullptr){
            return 0;
        }
        int left = getHeight(root->_left);
        int right = getHeight(root->_right);

        return left>right?left+1:right+1;
    }

    void RotateL(Node* parent){
        //右高左旋
        Node *cur = parent->_right;
        Node *curleft = cur->_left;

        parent->_right = curleft;
        if (curleft)
        {
            curleft->_parent = parent;
        }

        cur->_left = parent;

        Node *ppnode = parent->_parent;

        parent->_parent = cur;

        if (parent == _root)
        {
            _root = cur;
            cur->_parent = nullptr;
        }
        else
        {
            if (ppnode->_left == parent)
            {
                ppnode->_left = cur;
            }
            else
            {
                ppnode->_right = cur;
            }

            cur->_parent = ppnode;
        }

        parent->_bf = cur->_bf = 0;
    }

    void RotateR(Node *parent){ 
        //左高右旋
        Node *cur = parent->_left;
        Node *curright = cur->_right;

        parent->_left = curright;
        if (curright)
            curright->_parent = parent;

        Node *ppnode = parent->_parent;
        cur->_right = parent;
        parent->_parent = cur;

        if (ppnode == nullptr)
        {
            _root = cur;
            cur->_parent = nullptr;
        }
        else
        {
            if (ppnode->_left == parent)
            {
                ppnode->_left = cur;
            }
            else
            {
                ppnode->_right = cur;
            }

            cur->_parent = ppnode;
        }

        parent->_bf = cur->_bf = 0;
    }

    void RotateLR(Node *parent){
        Node *cur = parent->_left;
        Node *curright = cur->_right;
        int bf = curright->_bf;

        RotateL(parent->_left);
        RotateR(parent);

        if (bf == 0)
        {
            parent->_bf = 0;
            cur->_bf = 0;
            curright->_bf = 0;
        }
        else if (bf == -1)
        {
            parent->_bf = 1;
            cur->_bf = 0;
            curright->_bf = 0;
        }
        else if (bf == 1)
        {
            parent->_bf = 0;
            cur->_bf = -1;
            curright->_bf = 0;
        }
    }

    void RotateRL(Node *parent){
        Node *cur = parent->_right;
        Node *curleft = cur->_left;
        int bf = curleft->_bf;

        RotateR(parent->_right);
        RotateL(parent);

        if (bf == 0)
        {
            cur->_bf = 0;
            curleft->_bf = 0;
            parent->_bf = 0;
        }
        else if (bf == 1)
        {
            cur->_bf = 0;
            curleft->_bf = 0;
            parent->_bf = -1;
        }
        else if (bf == -1)
        {
            cur->_bf = 1;
            curleft->_bf = 0;
            parent->_bf = 0;
        }
        else
        {
            assert(false);
        }
    }

    void _InOrder(const Node *root)
    {
        if (root == nullptr)
            return;

        _InOrder(root->_left);

        std::cout << root->_kv.first << " : " << root->_kv.second << std::endl;

        _InOrder(root->_right);
    }

    private:

    Node *_root = nullptr;
};
